import os
import sys
import numpy as np
from PIL import Image
from typing import Dict, Any, Optional, Union, List
from pathlib import Path
import logging

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference


class SegmenterImpl:
    """
    分割器实现类
    
    提供图像分割功能，支持单张图片和批量分割推理。
    """
    
    def __init__(self, 
                 model_path: str, 
                 device: str = 'cpu',
                 class_names: Optional[List[str]] = None,
                 logger: Optional[logging.Logger] = None):
        """
        初始化分割器
        
        Args:
            model_path (str): 分割模型文件路径
            device (str): 推理设备，'cpu' 或 'cuda'
            class_names (List[str], optional): 类别名称列表，如果为None则自动从模型配置中获取
            logger (logging.Logger, optional): 日志记录器，如果提供则使用logger记录日志
        """
        self.model_path = model_path
        self.device = device
        self.class_names = class_names
        self.logger = logger
        
        # 创建推理器实例
        try:
            self.infer = create_inference(model_path, device=device, class_names=class_names)
            
            # 使用logger记录设备信息（如果提供了logger）
            if self.logger:
                self.logger.info(f"分割器初始化成功，使用设备: {device}")
            else:
                print(f"分割器初始化成功，使用设备: {device}")
            
            # 如果用户没有提供class_names，尝试从推理器中获取
            if self.class_names is None and hasattr(self.infer, 'class_names') and self.infer.class_names:
                self.class_names = self.infer.class_names
                if self.logger:
                    self.logger.info(f"自动获取类别名称: {self.class_names}")
                else:
                    print(f"自动获取类别名称: {self.class_names}")
            elif self.class_names is None:
                if self.logger:
                    self.logger.info("未找到类别名称，将使用默认类别名称")
                else:
                    print("未找到类别名称，将使用默认类别名称")
                
        except Exception as e:
            if self.logger:
                self.logger.error(f"初始化分割器失败: {str(e)}")
            raise RuntimeError(f"初始化分割器失败: {str(e)}")
    
    def segment_single(self, image: Union[str, Image.Image, np.ndarray], 
                      output_path: Optional[str] = None) -> Dict[str, Any]:
        """
        分割单张图片
        
        Args:
            image: 图像数据，可以是图像路径(str)、PIL Image对象、numpy数组、OpenCV BGR格式数组或单通道灰度图
            output_path (str, optional): 输出图片路径，如果提供则保存分割结果图片
            
        Returns:
            Dict[str, Any]: 分割结果，包含分割掩码、类别数量等信息
        """
        try:
            # 使用推理器进行单张图片推理
            result = self.infer.infer_single(image, return_prob=False)
            
            # 检查返回结果是否包含必要的字段
            if 'segmentation_mask' not in result:
                return {
                    'error': f"推理结果缺少segmentation_mask字段: {result}",
                    'success': False
                }
            
            # 如果提供了输出路径，可以在这里添加可视化功能（如果需要的话）
            if output_path:
                print(f"注意: 分割任务结果图片保存功能待实现")
            
            return {
                'segmentation_mask': result['segmentation_mask'],
                'num_classes': result.get('num_classes', None),
                'image_shape': result['segmentation_mask'].shape,
                'class_names': self.class_names if self.class_names else None,
                'success': True
            }
            
        except Exception as e:
            return {
                'error': str(e),
                'success': False
            }
    
    def segment_batch(self, input_dir: str) -> Dict[str, Any]:
        """
        批量分割目录中的图片
        
        Args:
            input_dir (str): 输入图片目录
            
        Returns:
            Dict[str, Any]: 批量分割结果，每张图片包含分割结果信息
        """
        try:
            # 使用推理器进行批量推理
            results = self.infer.infer_batch(input_dir)
            
            # 格式化结果
            formatted_results = {}
            for img_name, result in results.items():
                if 'error' not in result:
                    formatted_results[img_name] = {
                        'segmentation_mask': result['segmentation_mask'],
                        'num_classes': result.get('num_classes', None),
                        'image_shape': result['segmentation_mask'].shape,
                        'class_names': self.class_names if self.class_names else None,
                        'success': True
                    }
                else:
                    formatted_results[img_name] = {
                        'error': result['error'],
                        'success': False
                    }
            
            return formatted_results
            
        except Exception as e:
            return {
                'error': str(e),
                'success': False
            }
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        return self.infer.get_model_info()
    
    def validate_model(self) -> bool:
        """验证模型是否有效"""
        return self.infer.validate_model()
    
    def get_task_type(self) -> str:
        """获取任务类型"""
        return self.infer.get_task_type()

